#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json
from typing import List

from agentscope.agent import AgentBase
from agentscope.message import Msg
from agentscope.model import DashScopeChatModel
from agentscope.formatter import DashScopeChatFormatter
from agentscope.memory import InMemoryMemory
from pydantic import BaseModel, Field

from src.utils.env_config import aliyun_model_config

class Article(BaseModel):
    """标题"""
    title: str = Field(description="标题")
    content: str = Field(description="内容")
    relevance_score: int = Field(description="相关度得分")

class ArticleList(BaseModel):
    """文献检索结果"""
    retrieval_summary: List[Article] = Field(description="文献检索结果列表")

class LiteratureRetrieverAgent(AgentBase):
    """T2 - 文献检索智能体"""

    def __init__(self):
        super().__init__()
        self.name = "LiteratureRetrieverAgent"
        self.sys_prompt = """你是一个文献检索模拟智能体。你的任务是根据提供的关键词，生成一段模拟的文献检索结果摘要。
        要求：
        1. 基于关键词组合，构造合理的学术文献检索场景；
        2. 生成一段不少于50字的"检索结果摘要"，描述检索到的相关研究、论文数量、主要发现等；
        3. 内容需合理、连贯，体现专业性；
        4. 不需要真实调用数据库，仅做模拟；
        5. 输出格式必须为JSON，包含一个字段"retrieval_summary"，值为检索到的文章列表；
        6. 文章列表仅包含title、content、relevance_score字段，其中relevance_score为[0,100]内的整数
        7. 不要解释过程，不要添加额外内容。
        示例输出：
        {
            "retrieval_summary": [{
            "title": "基于深度学习的医学图像分析综述",
            "content": "本文综述了近年来深度学习应用进展...",
            "relevance_score": 95
            },
            {
            "title": "Transformer 在深度学习中的应用",
            "content": "Transformer 架构因其自注意力机制表现优异...",
            "relevance_score": 87
            }
        ]
        }
        """
        self.model = DashScopeChatModel(
            model_name=aliyun_model_config["model_name"],
            api_key=aliyun_model_config["api_key"],
            stream=False,
        )
        self.formatter = DashScopeChatFormatter()
        self.memory = InMemoryMemory()

    async def reply(self, msg: Msg) -> Msg:
        """处理文献检索请求"""
        await self.memory.add(msg)

        # 准备提示
        prompt = await self.formatter.format([
            Msg("system", self.sys_prompt, "system"),
            msg
        ])

        # 调用模型
        response = await self.model(prompt,structured_model=ArticleList)

        # 创建回复消息
        reply_msg = Msg(
            name=self.name,
            content=response.content,
            role="assistant"
        )

        await self.memory.add(reply_msg)
        return reply_msg

    async def retrieve_literature(self, keywords: List[str]) -> str:
        """检索文献"""
        user_msg = Msg(
            name="user",
            content=f"现在开始处理：\n关键词列表：{keywords}",
            role="user"
        )
        response_msg = await self.reply(user_msg)
        try:
            # 提取文本内容
            content_text = ""
            if isinstance(response_msg.content, list):
                for block in response_msg.content:
                    if hasattr(block, 'text'):
                        content_text += block.text
                    elif isinstance(block, dict) and 'text' in block:
                        content_text += block['text']
            else:
                content_text = str(response_msg.content)

            result = json.loads(content_text)
            return result.get("retrieval_summary", "")
        except json.JSONDecodeError:
            # 如果JSON解析失败，返回默认摘要
            print(f"----JSON解析失败：{response_msg.content}")
            return f"基于关键词{keywords}的文献检索结果显示，相关研究领域正在快速发展，具有重要的学术价值和应用前景。"

    async def observe(self, msg: Msg) -> None:
        """观察消息"""
        await self.memory.add(msg)

    async def handle_interrupt(self) -> Msg:
        """处理中断"""
        return Msg(
            name=self.name,
            content="文献检索任务被中断",
            role="assistant"
        )

